from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Path
from pydantic import BaseModel, Field, field_validator
import requests
import os
from dotenv import load_dotenv
from typing import Optional, List, Dict, Any
from fastapi.responses import JSONResponse
from enum import Enum
import base64

# 加载环境变量
load_dotenv()

app = FastAPI(title="生成魔法工坊")

# 定义照片风格枚举
class ImageStyle(str, Enum):
    cartoon = "卡通风格"
    oil_painting = "油画风格"
    sketch = "素描风格"
    ink_wash = "水墨风格"

# 定义图片尺寸枚举
class ImageSizeEnum(str, Enum):
    _1024x1024 = "1024x1024"
    _960x1280 = "960x1280"
    _768x1024 = "768x1024"
    _720x1440 = "720x1440"
    _720x1280 = "720x1280"

# 定义3D模型艺术风格枚举
class ArtStyleEnum(str, Enum):
    realistic = "realistic"   # 写实风格，默认值
    sculpture = "sculpture"   # 雕塑风格，使用此风格时enable_pbr会自动设置为false

# 基础响应模型
class ImageResponse(BaseModel):
    url: str

class TimingsResponse(BaseModel):
    inference: float

class SuccessResponse(BaseModel):
    images: List[ImageResponse]
    timings: TimingsResponse
    seed: int

class ErrorResponse(BaseModel):
    code: Optional[int] = None
    message: str
    data: Optional[str] = None

# 3D模型相关模型
class ModelUrls(BaseModel):
    glb: str
    fbx: str
    obj: str
    mtl: str
    usdz: str

class TextureUrls(BaseModel):
    base_color: str

class TaskError(BaseModel):
    message: str

class TextTo3DTaskResponse(BaseModel):
    id: str
    model_urls: ModelUrls
    thumbnail_url: str
    prompt: str
    art_style: str
    progress: int
    started_at: int
    created_at: int
    finished_at: int
    status: str
    texture_urls: List[TextureUrls]
    preceding_tasks: int
    task_error: TaskError

class TextTo3DResponse(BaseModel):
    result: str

# 图片生成请求模型
class ImageGenerationRequest(BaseModel):
    prompt: str = Field(..., description="图片生成的文本描述")
    image_size: ImageSizeEnum = Field(description="图片尺寸")
    num_inference_steps: int = Field(
        default=20,
        ge=1,
        le=100,
        description="推理步数（范围：1 ≤ x ≤ 100）"
    )
    guidance_scale: float = Field(
        default=7.5,
        ge=0,
        le=20,
        description="引导强度（范围：0 ≤ x ≤ 20）"
    )
    negative_prompt: Optional[str] = Field(default="", description="负面提示词")
    seed: Optional[int] = Field(
        default=4999999999,
        description="随机种子",
    )

    class Config:
        json_schema_extra = {
            "example": {
                "prompt": "an island near sea, with seagulls, moon shining over the sea",
                "image_size": ImageSizeEnum._1024x1024.value,
                "num_inference_steps": 20,
                "guidance_scale": 7.5,
                "negative_prompt": ""
            }
        }

# 3D模型生成请求模型
class TextTo3DRequest(BaseModel):
    prompt: str = Field(..., description="3D模型描述文本，最多600字符")
    art_style: ArtStyleEnum = Field(
        default=ArtStyleEnum.realistic,
        description="艺术风格：realistic(写实风格，默认值) 或 sculpture(雕塑风格)"
    )

    class Config:
        json_schema_extra = {
            "example": {
                "prompt": "一头牛",
                "art_style": "realistic"
            }
        }

class Refine3DRequest(BaseModel):
    preview_task_id: str = Field(..., description="对应的预览任务id，从预览任务接口获取")
    enable_pbr: bool = Field(False, description="是否生成PBR贴图，当使用雕塑风格时应设置为false")
    texture_prompt: Optional[str] = Field(None, exclude=True)
    texture_image_url: Optional[str] = Field(None, exclude=True)

    class Config:
        json_schema_extra = {
            "example": {
                "preview_task_id": "01973942-5a45-7f96-bb78-1742c0d99e19",
                "enable_pbr": True
            }
        }

    @field_validator('texture_prompt', 'texture_image_url')
    @classmethod
    def validate_texture_params(cls, v, info):
        values = info.data
        if 'texture_prompt' in values and 'texture_image_url' in values:
            if values['texture_prompt'] and values['texture_image_url']:
                raise ValueError("texture_prompt 和 texture_image_url 不能同时提供")
        return v

# 异常处理器
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    return JSONResponse(
        status_code=exc.status_code,
        content=exc.detail
    )

# 图片生成接口
@app.post("/generate-image", response_model=SuccessResponse)
async def generate_image(request: ImageGenerationRequest):
    # image_size直接使用request.image_size.value获取字符串值
    try:
        payload = {
            "model": "Kwai-Kolors/Kolors",
            "prompt": request.prompt,
            "negative_prompt": request.negative_prompt,
            "image_size": request.image_size.value, # 使用枚举的字符串值
            "batch_size": 1,
            "seed": request.seed,
            "num_inference_steps": request.num_inference_steps,
            "guidance_scale": request.guidance_scale
        }

        headers = {
            "Authorization": f"Bearer {os.getenv('API_KEY')}",
            "Content-Type": "application/json"
        }

        response = requests.post(
            os.getenv('API_URL'),
            json=payload,
            headers=headers
        )
        
        if response.status_code == 200:
            return response.json()
        elif response.status_code == 400:
            error_data = response.json()
            raise HTTPException(
                status_code=400,
                detail={
                    "code": error_data.get("code"),
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 401:
            raise HTTPException(status_code=401, detail="Invalid token")
        elif response.status_code == 404:
            raise HTTPException(status_code=404, detail="404 page not found")
        elif response.status_code == 429:
            error_data = response.json()
            raise HTTPException(
                status_code=429,
                detail={
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 503:
            error_data = response.json()
            raise HTTPException(
                status_code=503,
                detail={
                    "code": error_data.get("code"),
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 504:
            raise HTTPException(status_code=504, detail=response.text)
        else:
            raise HTTPException(
                status_code=response.status_code,
                detail=f"Unexpected error: {response.text}"
            )
            
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Request failed: {str(e)}")
    except Exception as e:
        print(f"An unexpected error occurred in /generate-image: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 图片风格转换接口
@app.post("/convert-style", response_model=SuccessResponse)
async def convert_style(
    image: UploadFile = File(..., description="需要上传的照片"),
    style: ImageStyle = Form(..., description="选择要转换的照片风格"),
    image_size: ImageSizeEnum = Form(description="输出图片尺寸"),
    num_inference_steps: int = Form(default=20, ge=1, le=100, description="推理步数（范围：1 ≤ x ≤ 100）"),
    guidance_scale: float = Form(default=7.5, ge=0, le=20, description="引导强度（范围：0 ≤ x ≤ 20）"),
    negative_prompt: Optional[str] = Form(default="", description="负面提示词"),
):
    try:
        image_data = await image.read()
        content_type = image.content_type
        if content_type and content_type.startswith("image/"):
            base64_prefix = f"data:{content_type};base64,"
        else:
            base64_prefix = "data:image/png;base64,"
            print(f"Warning: Could not determine image type from content_type: {content_type}. Using default png prefix.")

        base64_image = base64.b64encode(image_data).decode('utf-8')
        base64_image_with_prefix = f"{base64_prefix}{base64_image}"

        style_prompts = {
            ImageStyle.cartoon: f"a photo of the image in cartoon style illustration",
            ImageStyle.oil_painting: f"a photo of the image as an oil painting",
            ImageStyle.sketch: f"a photo of the image as a sketch",
            ImageStyle.ink_wash: f"a photo of the image as an ink wash painting"
        }
        conversion_prompt = style_prompts.get(style, f"a photo of the image in {style.value} style")

        payload = {
            "model": "Kwai-Kolors/Kolors",
            "prompt": conversion_prompt,
            "image": base64_image_with_prefix,
            "image_size": image_size.value, # 使用枚举的字符串值
            "batch_size": 1,
            "seed": 4999999999,
            "num_inference_steps": num_inference_steps,
            "guidance_scale": guidance_scale,
            "negative_prompt": negative_prompt
        }

        api_url = os.getenv('API_URL')
        api_key = os.getenv('API_KEY')

        if not api_url or not api_key:
             missing_config = []
             if not api_url: missing_config.append('API_URL')
             if not api_key: missing_config.append('API_KEY')
             raise HTTPException(status_code=500, detail=f"Missing API configuration: {', '.join(missing_config)}")

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        response = requests.post(
            api_url,
            json=payload,
            headers=headers
        )

        if response.status_code == 200:
            return response.json()
        elif response.status_code == 400:
            error_data = response.json()
            raise HTTPException(
                status_code=400,
                detail={
                    "code": error_data.get("code"),
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 401:
            raise HTTPException(status_code=401, detail="Invalid token")
        elif response.status_code == 404:
            raise HTTPException(status_code=404, detail="404 page not found")
        elif response.status_code == 429:
            error_data = response.json()
            raise HTTPException(
                status_code=429,
                detail={
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 503:
            error_data = response.json()
            raise HTTPException(
                status_code=503,
                detail={
                    "code": error_data.get("code"),
                    "message": error_data.get("message"),
                    "data": error_data.get("data")
                }
            )
        elif response.status_code == 504:
            raise HTTPException(status_code=504, detail=response.text)
        else:
            raise HTTPException(
                status_code=response.status_code,
                detail=f"Unexpected error: {response.text}"
            )

    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"API request failed: {str(e)}")
    except Exception as e:
        print(f"An unexpected error occurred in /convert-style: {e}")
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")

# 3D模型生成接口
@app.post("/text-to-3d", response_model=TextTo3DResponse)
async def text_to_3d(req: TextTo3DRequest):
    """
    创建3D模型预览任务
    
    参数说明：
    - prompt: 必填，3D模型描述文本，最多600字符
    - art_style: 可选，艺术风格
      - realistic: 写实风格（默认值）
      - sculpture: 雕塑风格（使用此风格时enable_pbr会自动设置为false）
    
    返回说明：
    - result: 任务ID，用于后续查询任务状态或创建精细化任务
    """
    enable_pbr = False if req.art_style == ArtStyleEnum.sculpture else True
    payload = {
        "mode": "preview",
        "prompt": req.prompt,
        "art_style": req.art_style,
        "ai_model": "meshy-4",
        "topology": "quad",
        "target_polycount": 30000,
        "should_remesh": True,
        "symmetry_mode": "auto",
        "enable_pbr": enable_pbr
    }
    headers = {
        "Authorization": f"Bearer {os.getenv('MESHY_API_KEY')}"
    }
    print("预览任务请求参数:", payload)
    try:
        response = requests.post(os.getenv('MESHY_API_URL'), headers=headers, json=payload)
        print("API响应状态码:", response.status_code)
        print("API响应内容:", response.text)
        return response.json()
    except Exception as e:
        print("异常:", str(e))
        return {"result": str(e)}

@app.post("/text-to-3d/refine")
async def refine_3d(req: Refine3DRequest):
    """
    创建3D模型精细化任务
    
    参数说明：
    - preview_task_id: 必填，从预览任务接口获取的任务ID
    - enable_pbr: 可选，默认为false，是否生成PBR贴图
    
    返回说明：
    - result: 任务ID，用于后续查询任务状态
    """
    payload = {
        "mode": "refine",
        "preview_task_id": req.preview_task_id,
        "enable_pbr": req.enable_pbr
    }
    
    if req.texture_prompt is not None:
        payload["texture_prompt"] = req.texture_prompt
    if req.texture_image_url is not None:
        payload["texture_image_url"] = req.texture_image_url

    headers = {
        "Authorization": f"Bearer {os.getenv('MESHY_API_KEY')}"
    }
    print("精细化任务请求参数:", payload)
    try:
        response = requests.post(os.getenv('MESHY_API_URL'), headers=headers, json=payload)
        print("API响应状态码:", response.status_code)
        print("API响应内容:", response.text)
        return response.json()
    except Exception as e:
        print("异常:", str(e))
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/text-to-3d/{task_id}")
async def get_text_to_3d_result(
    task_id: str = Path(..., description="文本生成3D任务的唯一ID")
):
    """
    查询3D模型任务状态
    
    参数说明：
    - task_id: 必填，任务ID，从预览任务或精细化任务接口返回的result字段获取
    
    返回说明：
    - status: 任务状态（PENDING/PROCESSING/SUCCEEDED/FAILED）
    - progress: 任务进度（0-100）
    - model_urls: 模型下载链接
    - texture_urls: 贴图下载链接（如果启用了PBR）
    """
    headers = {
        "Authorization": f"Bearer {os.getenv('MESHY_API_KEY')}"
    }
    url = f"https://api.meshy.ai/openapi/v2/text-to-3d/{task_id}"
    print("查询任务ID:", task_id)
    try:
        response = requests.get(url, headers=headers)
        print("API响应状态码:", response.status_code)
        print("API响应内容:", response.text)
        return response.json()
    except Exception as e:
        print("异常:", str(e))
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/image-to-3d/upload")
async def image_to_3d_upload(
    file: UploadFile = File(..., description="上传的图片文件，支持jpg、jpeg、png格式")
):
    """
    通过上传图片创建3D模型任务
    
    参数说明：
    - file: 必填，上传的图片文件（支持jpg、jpeg、png格式）
    
    返回说明：
    - result: 任务ID，用于后续查询任务状态
    """
    if not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="只支持图片文件（jpg、jpeg、png）")
    
    file_content = await file.read()
    base64_content = base64.b64encode(file_content).decode('utf-8')
    data_uri = f"data:{file.content_type};base64,{base64_content}"
    
    payload = {
        "image_url": data_uri,
        "ai_model": "meshy-4",
        "topology": "quad",
        "target_polycount": 30000,
        "symmetry_mode": "auto",
        "should_remesh": True,
        "should_texture": True,
        "enable_pbr": True
    }

    headers = {
        "Authorization": f"Bearer {os.getenv('MESHY_API_KEY')}"
    }
    url = "https://api.meshy.ai/openapi/v1/image-to-3d"
    print("图像转3D任务请求参数:", {**payload, "image_url": "data:image/...base64..."})
    try:
        response = requests.post(url, headers=headers, json=payload)
        print("API响应状态码:", response.status_code)
        print("API响应内容:", response.text)
        return response.json()
    except Exception as e:
        print("异常:", str(e))
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/image-to-3d/{task_id}")
async def get_image_to_3d_result(
    task_id: str = Path(..., description="图像转3D任务的唯一ID")
):
    """
    查询图像转3D任务状态
    
    参数说明：
    - task_id: 必填，任务ID，从图像转3D任务接口返回的result字段获取
    
    返回说明：
    - status: 任务状态（PENDING/PROCESSING/SUCCEEDED/FAILED）
    - progress: 任务进度（0-100）
    - model_urls: 模型下载链接
    - texture_urls: 贴图下载链接（如果启用了贴图）
    """
    headers = {
        "Authorization": f"Bearer {os.getenv('MESHY_API_KEY')}"
    }
    url = f"https://api.meshy.ai/openapi/v1/image-to-3d/{task_id}"
    print("查询任务ID:", task_id)
    try:
        response = requests.get(url, headers=headers)
        print("API响应状态码:", response.status_code)
        print("API响应内容:", response.text)
        return response.json()
    except Exception as e:
        print("异常:", str(e))
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=1189)